from typing import Optional
from chernc.treediff.gumtree.actions.chawathe_script_generator import ChawatheScriptGenerator
from chernc.treediff.gumtree.actions.edit_script import EditScript
from chernc.treediff.gumtree.actions.edit_script_generator import EditScriptGenerator
from chernc.treediff.gumtree.actions.model.delete import Delete
from chernc.treediff.gumtree.actions.model.insert import Insert
from chernc.treediff.gumtree.actions.model.tree_delete import TreeDelete
from chernc.treediff.gumtree.actions.model.tree_insert import TreeInsert
from chernc.treediff.gumtree.matchers.mapping_store import MappingStore


class SimplifiedChawatheScriptGenerator(EditScriptGenerator):
    def compute_actions(self, mappings: MappingStore) -> EditScript:
        actions = ChawatheScriptGenerator().compute_actions(mappings)
        return self.simplify(actions)

    @staticmethod
    def simplify(actions: 'EditScript') -> 'EditScript':
        added_trees = {}
        deleted_trees = {}

        # Categorize the actions
        for a in actions:
            if isinstance(a, Insert):
                added_trees[a.get_node()] = a
            elif isinstance(a, Delete):
                deleted_trees[a.get_node()] = a

        # Simplify the actions related to added trees
        for t in list(added_trees.keys()):
            if t.get_parent() in added_trees and all(descendant in added_trees for descendant in t.get_parent().get_descendants()):
                actions.remove(added_trees[t])
            else:
                if len(t.get_children()) > 0 and all(descendant in added_trees for descendant in t.get_descendants()):
                    original_action = added_trees[t]
                    ti = TreeInsert(original_action.get_node(), original_action.get_parent(), original_action.get_position())
                    index = actions.last_index_of(original_action)
                    actions.add_at(index, ti)
                    actions.remove_at(index + 1)

        # Simplify the actions related to deleted trees
        for t in list(deleted_trees.keys()):
            if t.get_parent() in deleted_trees and all(descendant in deleted_trees for descendant in t.get_parent().get_descendants()):
                actions.remove(deleted_trees[t])
            else:
                if len(t.get_children()) > 0 and all(descendant in deleted_trees for descendant in t.get_descendants()):
                    original_action = deleted_trees[t]
                    ti = TreeDelete(original_action.get_node())
                    index = actions.last_index_of(original_action)
                    actions.add_at(index, ti)
                    actions.remove_at(index + 1)

        return actions
